from datatable import DataTable
import zipfile
import urllib
import StringIO

_valid_fields = set([
        # http://www.transtats.bts.gov/DL_SelectFields.asp?Table_ID=236
        'YEAR', #
        'QUARTER', # 1=Jan-Mar
        'MONTH', # 1=Jan
        'DAY_OF_MONTH', #
        'DAY_OF_WEEK', # 1=Monday
        'FL_DATE', # yyyymmdd
        'UNIQUE_CARRIER', #
        'AIRLINE_ID', # assigned by DOT
        'CARRIER', # assigned by IATA
        'TAIL_NUM', #
        'FL_NUM', #
        'ORIGIN', 'ORIGIN_CITY_NAME','ORIGIN_STATE_ABR','ORIGIN_STATE_FIPS','ORIGIN_STATE_NM','ORIGIN_WAC', #
        'DEST', 'DEST_CITY_NAME','DEST_STATE_ABR','DEST_STATE_FIPS','DEST_STATE_NM','DEST_WAC', #
        'CRS_DEP_TIME', 'DEP_TIME', 'DEP_DELAY', 'DEP_DELAY_NEW', 'DEP_DEL15', 'DEP_DELAY_GROUP', 'DEP_TIME_BLK', #
        'TAXI_OUT','WHEELS_OFF','WHEELS_ON','TAXI_IN', #
        'CRS_ARR_TIME', 'ARR_TIME', 'ARR_DELAY', 'ARR_DELAY_NEW', 'ARR_DEL15', 'ARR_DELAY_GROUP', 'ARR_TIME_BLK', #
        'CANCELLED','CANCELLATION_CODE', #
        'DIVERTED', #
        'CRS_ELAPSED_TIME', 'ACTUAL_ELAPSED_TIME', 'AIR_TIME', #
        'FLIGHTS', #
        'DISTANCE','DISTANCE_GROUP', #
        'CARRIER_DELAY','WEATHER_DELAY','NAS_DELAY','SECURITY_DELAY','LATE_AIRCRAFT_DELAY', #
        'FIRST_DEP_TIME', #
        'TOTAL_ADD_GTIME', #
        'LONGEST_ADD_GTIME', #
        'DIV_AIRPORT_LANDINGS', #
        'DIV_REACHED_DEST', #
        'DIV_ACTUAL_ELAPSED_TIME', #
        'DIV_ARR_DELAY', #
        'DIV_DISTANCE', #
        'DIV1_AIRPORT','DIV1_WHEELS_ON','DIV1_TOTAL_GTIME','DIV1_LONGEST_GTIME','DIV1_WHEELS_OFF','DIV1_TAIL_NUM', #
        'DIV2_AIRPORT','DIV2_WHEELS_ON','DIV2_TOTAL_GTIME','DIV2_LONGEST_GTIME','DIV2_WHEELS_OFF','DIV2_TAIL_NUM', #
        'DIV3_AIRPORT','DIV3_WHEELS_ON','DIV3_TOTAL_GTIME','DIV3_LONGEST_GTIME','DIV3_WHEELS_OFF','DIV3_TAIL_NUM', #
        'DIV4_AIRPORT','DIV4_WHEELS_ON','DIV4_TOTAL_GTIME','DIV4_LONGEST_GTIME','DIV4_WHEELS_OFF','DIV4_TAIL_NUM', #
        'DIV5_AIRPORT','DIV5_WHEELS_ON','DIV5_TOTAL_GTIME','DIV5_LONGEST_GTIME','DIV5_WHEELS_OFF','DIV5_TAIL_NUM'  #
    ])

def from_query(query_dict,cols=[
        'YEAR','MONTH','DAY_OF_MONTH',
        'DAY_OF_WEEK',
        'UNIQUE_CARRIER',
        'ORIGIN','ORIGIN_WAC',
        'DEST','DEST_WAC',
        'ARR_DELAY','CANCELLED','DIVERTED'
    ]):
    for col in cols:
        assert col in _valid_fields
    for key in query_dict.keys():
        assert key in _valid_fields

    sql_select = ','.join(cols)
    sql_where = ' AND '.join(
        map(lambda key: key + '=' + query_dict[key], query_dict)
    )
    sql = 'SELECT ' + sql_select + ' FROM T_ONTIME WHERE ' + sql_where

    url = 'http://www.transtats.bts.gov/DownLoad_Table.asp'
    params = {'sqlstr':sql}
    response = urllib.urlopen(url,urllib.urlencode(params))
    assert response.info().gettype() == 'application/x-zip-compressed'

    response_file = StringIO.StringIO(response.read())
    zip_file = zipfile.ZipFile(response_file, 'r')
    assert len(zip_file.namelist()) == 1

    file_name = zip_file.namelist()[0]
    return DataTable(
        zip_file.read(file_name)
    )

def from_year_month(year,month):
    return from_query({'YEAR':str(year),'MONTH':str(month)})

# I would have made these methods on FlightDataTable,
# except that summarize returns instances of the superclass,
# so they couldn't be chained. What's the right idiom for this?

def summarize_by_month(table):
    group_cols = ['YEAR','MONTH','UNIQUE_CARRIER','ORIGIN','DEST']
    summary_cols = {
            'COUNT':
                lambda summary_row, data_row:
                    summary_row['COUNT'] + 1,
            'COUNT_DELAYED':
                lambda summary_row, data_row:
                    summary_row['COUNT_DELAYED'] +
                        int(
                            float(data_row['ARR_DELAY'] or '0')>15 or
                            float(data_row['DIVERTED'] or '0') or
                            float(data_row['CANCELLED'] or '0')
                        )
        }
    return table.summarize(group_cols,summary_cols)

def insert_city_pair(table):
    def concat_sorted(orig,dest):
        list = [orig,dest]
        list.sort()
        return '-'.join(list)
    group_cols = ['YEAR','MONTH','UNIQUE_CARRIER','ORIGIN','DEST','COUNT','COUNT_DELAYED']
    summary_cols = {
        'CITY_PAIR':
            lambda summary_row, data_row:
                concat_sorted(data_row['ORIGIN'],data_row['DEST'])
        }
    return table.summarize(group_cols,summary_cols)

def summarize_by_city_pair(table):
    group_cols = ['YEAR','MONTH','UNIQUE_CARRIER','CITY_PAIR']
    summary_cols = {
            'COUNT':
                lambda summary_row, data_row:
                    int(summary_row['COUNT']) + int(data_row['COUNT']),
            'COUNT_DELAYED':
                lambda summary_row, data_row:
                    int(summary_row['COUNT_DELAYED']) + int(data_row['COUNT_DELAYED'])
        }
    return table.summarize(group_cols,summary_cols)

def summarize_by_month_city_pair(table):
    table = summarize_by_month(table)
    table = insert_city_pair(table)
    table = summarize_by_city_pair(table)
    return table
